#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
from jm.library.data_helper import filter_data, date_to_str
from jm.library.likelihood import Likelihood
from matplotlib import rcParams
import pandas as pd
import warnings
warnings.filterwarnings('ignore')

rcParams['figure.figsize'] = (13.0, 7.0)
rcParams['lines.linewidth'] = 3
rcParams['font.size'] = 18

# In[2]:

df_status = pd.read_csv('data/jesus_maria_status_deidentified.csv')

# In[3]:


def get_cols_over_time(col_fmt='payments_by_{}', starting=0):
    return [col_fmt.format(date_to_str(d)) for d in Likelihood.event_dates[starting:]]


def plot_over_time(df, col_fmt='payments_by_{}', func=np.sum, ax=None, starting=0,  linestyle='solid'):
    cols = get_cols_over_time(col_fmt, starting)
    this_df = df[cols]
    this_df.columns = Likelihood.event_dates[starting:]
    ax = this_df.apply(func, axis=0).plot(ax=ax,  linestyle=linestyle)
    return ax


def col_d(col_fmt, d):
    return col_fmt.format(date_to_str(d))


def generate_new_vars(plot_df):
    plot_df = plot_df.copy()
    for d in Likelihood.event_dates:
        plot_df.loc[plot_df[col_d('rec1_by_{}',d)]==1, 
                    col_d('rec1_fwdiff_by_{}',d)] = plot_df[col_d('relative_fwdiff_payments_by_{}',d)]/4
        plot_df.loc[plot_df[col_d('valor_by_{}',d)]==1, 
                    col_d('valor_fwdiff_by_{}',d)] = plot_df[col_d('relative_fwdiff_payments_by_{}',d)]/4
        plot_df.loc[plot_df[col_d('rec1valor_by_{}',d)]==1, 
                    col_d('rec1valor_fwdiff_by_{}',d)] = plot_df[col_d('relative_fwdiff_payments_by_{}',d)]/4
        plot_df.loc[plot_df[col_d('N_by_{}',d)]==1, 
                    col_d('N_fwdiff_by_{}',d)] = plot_df[col_d('relative_fwdiff_payments_by_{}',d)]/4

        plot_df.loc[(plot_df[col_d('rec1_by_{}',d)]==1) & (plot_df[col_d('fwdiff_payments_by_{}',d)]>0), 
                    col_d('pos_rec1_fwdiff_by_{}',d)] = plot_df[col_d('relative_fwdiff_payments_by_{}',d)]
        plot_df.loc[(plot_df[col_d('valor_by_{}',d)]==1) & (plot_df[col_d('fwdiff_payments_by_{}',d)]>0), 
                    col_d('pos_valor_fwdiff_by_{}',d)] = plot_df[col_d('relative_fwdiff_payments_by_{}',d)]
        plot_df.loc[(plot_df[col_d('rec1valor_by_{}',d)]==1) & (plot_df[col_d('fwdiff_payments_by_{}',d)]>0), 
                    col_d('pos_rec1valor_fwdiff_by_{}',d)] = plot_df[col_d('relative_fwdiff_payments_by_{}',d)]
        plot_df.loc[(plot_df[col_d('N_by_{}',d)]==1) & (plot_df[col_d('fwdiff_payments_by_{}',d)]>0), 
                    col_d('pos_N_fwdiff_by_{}',d)] = plot_df[col_d('relative_fwdiff_payments_by_{}',d)]

        plot_df.loc[plot_df[col_d('rec1_by_{}',d)]==1, 
                    col_d('ind_rec1_fwdiff_by_{}',d)] = plot_df[col_d('fwdiff_payments_by_{}',d)] > 0
        plot_df.loc[plot_df[col_d('valor_by_{}',d)]==1, 
                    col_d('ind_valor_fwdiff_by_{}',d)] = plot_df[col_d('fwdiff_payments_by_{}',d)] > 0
        plot_df.loc[plot_df[col_d('rec1valor_by_{}',d)]==1, 
                    col_d('ind_rec1valor_fwdiff_by_{}',d)] = plot_df[col_d('fwdiff_payments_by_{}',d)] > 0
        plot_df.loc[plot_df['N_by_{}'.format(date_to_str(d))]==1, 
                    col_d('ind_N_fwdiff_by_{}',d)] = plot_df[col_d('fwdiff_payments_by_{}',d)] > 0
    return plot_df


payment_cols = get_cols_over_time()
priority_cols = get_cols_over_time('priority_by_{}')
action_cols = get_cols_over_time('action_by_{}')
action_cols_rec1 = get_cols_over_time('rec1_by_{}')
action_cols_rec1valor = get_cols_over_time('rec1valor_by_{}')
action_cols_N = get_cols_over_time('N_by_{}')
action_cols_valor = get_cols_over_time('valor_by_{}')
action_no_medida_cols = get_cols_over_time('action_no_medida_by_{}')
relative_payment_cols = ['relative_' + c for c in payment_cols]
fwdiff_payment_cols = ['fwdiff_' + p for p in payment_cols]
relative_fwdiff_payment_cols = ['relative_' + c for c in fwdiff_payment_cols]


df_status[relative_payment_cols] = df_status[payment_cols].divide(
                df_status['Q1_total_due'], axis=0)


df_status[relative_fwdiff_payment_cols] = df_status[fwdiff_payment_cols].divide(
    df_status['Q1_total_due'], axis=0)


df_status[action_cols_rec1] = df_status[action_cols]=='rec1'
df_status[action_cols_valor] = df_status[action_cols]=='valor'
df_status[action_cols_rec1valor] = (df_status[action_cols]=='rec1') | (df_status[action_cols]=='valor')
df_status[action_cols_N] = df_status[action_cols]=='N'


df_treatment = filter_data(df_status, assignment_to_treatment=1)
df_control = filter_data(df_status, assignment_to_treatment=0)


plot_df = generate_new_vars(df_control)
plt.clf()
ax = plot_over_time(plot_df, col_fmt='valor_fwdiff_by_{}', func=np.mean, starting=0)
plot_over_time(plot_df, col_fmt='N_fwdiff_by_{}', func=np.mean, ax=ax, starting=0,  linestyle='dashed')
plt.legend(['latest action: notification', 'latest action: none'])
plt.ylabel('Weekly Payment / Annualized Tax Due')
plt.tight_layout()
plt.savefig('figs/figOB4.pdf')


plot_df = generate_new_vars(df_treatment)
plt.clf()
ax = plot_over_time(plot_df, col_fmt='valor_fwdiff_by_{}', func=np.mean, starting=0)
plot_over_time(plot_df, col_fmt='N_fwdiff_by_{}', func=np.mean, ax=ax, starting=0,  linestyle='dashed')
plt.legend(['latest action: notification', 'latest action: none'])
plt.ylabel('Weekly Payment / Annualized Tax Due')
plt.tight_layout()
plt.savefig('figs/figOB5.pdf')

